import hashlib

# AES核心参数定义
AES_BLOCK_SIZE = 16
AES_ROUNDS = 10

# S盒
S_BOX = [
    0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
    0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
    0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
    0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
    0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
    0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
    0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
    0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
    0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
    0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
    0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
    0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
    0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
    0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
    0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
    0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
]

# 轮常数
R_CON = [0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36]

class AESCMAC:
    """AES-CMAC算法工具类 (RFC4493标准实现)"""
    
    @staticmethod
    def xor_bytes(a, b):
        """字节数组异或操作"""
        return bytes(x ^ y for x, y in zip(a, b))
    
    @staticmethod
    def left_shift(block):
        """左移1位操作"""
        output = bytearray(block)
        carry = 0
        for i in reversed(range(len(output))):
            new_carry = (output[i] & 0x80) >> 7
            output[i] = ((output[i] << 1) | carry) & 0xFF
            carry = new_carry
        return bytes(output)
    
    @staticmethod
    def sub_bytes(state):
        """字节替换操作（使用S盒）"""
        return bytes(S_BOX[byte] for byte in state)
    
    @staticmethod
    def shift_rows(state):
        """行移位操作"""
        state_bytes = bytearray(state)
        # 第一行循环左移1位
        state_bytes[1], state_bytes[5], state_bytes[9], state_bytes[13] = state_bytes[5], state_bytes[9], state_bytes[13], state_bytes[1]
        # 第二行循环左移2位
        state_bytes[2], state_bytes[10] = state_bytes[10], state_bytes[2]
        state_bytes[6], state_bytes[14] = state_bytes[14], state_bytes[6]
        # 第三行循环左移3位（正确实现：d,a,b,c）
        state_bytes[3], state_bytes[7], state_bytes[11], state_bytes[15] = state_bytes[15], state_bytes[3], state_bytes[7], state_bytes[11]
        return bytes(state_bytes)
    
    @staticmethod
    def xtime(x):
        """列混合中的字节乘法（x*2 mod 0x1b）"""
        return ((x << 1) ^ 0x1b) & 0xFF if x & 0x80 else (x << 1) & 0xFF
    
    @staticmethod
    def mix_columns(state):
        """列混合操作"""
        state_bytes = bytearray(state)
        for i in range(0, 16, 4):
            a = state_bytes[i]
            b = state_bytes[i+1]
            c = state_bytes[i+2]
            d = state_bytes[i+3]
            
            state_bytes[i] = AESCMAC.xtime(a) ^ AESCMAC.xtime(b) ^ b ^ c ^ d
            state_bytes[i+1] = a ^ AESCMAC.xtime(b) ^ AESCMAC.xtime(c) ^ c ^ d
            state_bytes[i+2] = a ^ b ^ AESCMAC.xtime(c) ^ AESCMAC.xtime(d) ^ d
            state_bytes[i+3] = AESCMAC.xtime(a) ^ a ^ b ^ c ^ AESCMAC.xtime(d)
        return bytes(state_bytes)
    
    @staticmethod
    def add_round_key(state, round_key):
        """轮密钥加操作"""
        return AESCMAC.xor_bytes(state, round_key)
    
    @staticmethod
    def key_expansion(key):
        """密钥扩展（生成轮密钥）"""
        key_size = len(key)
        if key_size not in (16, 24, 32):
            raise ValueError("仅支持128/192/256位密钥")
        rounds = 10 if key_size == 16 else 12 if key_size == 24 else 14
        # 计算轮密钥总长度（轮数+1）* 块大小（16字节）
        total_key_size = (rounds + 1) * 16
        # 初始化轮密钥数组，前key_size字节为原始密钥，后续填充0
        round_keys = bytearray(total_key_size)
        round_keys[:len(key)] = key
        
        for i in range(4, (rounds + 1) * 4):
            temp = round_keys[(i-1)*4 : i*4]
            if i % 4 == 0:
                # 字循环 + 字节替换 + 轮常数异或
                temp = temp[1:] + temp[:1]  # 字循环
                temp = bytes(S_BOX[b] for b in temp)  # 字节替换
                temp = bytes([temp[0] ^ R_CON[i//4]]) + temp[1:]
            
            for j in range(4):
                round_keys[i*4 + j] = round_keys[(i-4)*4 + j] ^ temp[j]
        
        return bytes(round_keys)
    
    @staticmethod
    def aes_encrypt(plaintext, key):
        """AES加密核心函数"""
        state = bytearray(plaintext)
        round_keys = AESCMAC.key_expansion(key)
        num_rounds = AES_ROUNDS if len(key) == 16 else 12 if len(key) == 24 else 14
        
        # 初始轮密钥加
        state = AESCMAC.add_round_key(state, round_keys[:16])

        # 主轮次加密
        for round_num in range(1, num_rounds):
            state = AESCMAC.sub_bytes(state)
            state = AESCMAC.shift_rows(state)
            state = AESCMAC.mix_columns(state)
            state = AESCMAC.add_round_key(state, round_keys[round_num*16 : (round_num+1)*16])
        
        # 最后一轮（无列混合）
        state = AESCMAC.sub_bytes(state)
        state = AESCMAC.shift_rows(state)
        state = AESCMAC.add_round_key(state, round_keys[num_rounds*16 : (num_rounds+1)*16])
        
        return bytes(state)
    
    @staticmethod
    def generate_subkeys(key):
        """生成CMAC子密钥K1和K2"""
        # 使用自定义AES加密全零块
        zero_block = bytes(16)
        L = AESCMAC.aes_encrypt(zero_block, key)
        
        # 生成K1
        K1 = AESCMAC.left_shift(L)
        if L[0] & 0x80:
            K1 = AESCMAC.xor_bytes(K1, bytes.fromhex('00000000000000000000000000000087'))
            
        # 生成K2
        K2 = AESCMAC.left_shift(K1)
        if K1[0] & 0x80:
            K2 = AESCMAC.xor_bytes(K2, bytes.fromhex('00000000000000000000000000000087'))
            
        return K1, K2

    @staticmethod
    def calculate(data, key = [0xF4, 0x6B, 0x20, 0xEA, 0xFE, 0x88, 0x98, 0x58, 0x27, 0xBB, 0x7C, 0x9A, 0xC7, 0x32, 0xE4, 0x66]):
        """
        计算AES-CMAC值
        :param data: 输入数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :return: CMAC结果(16字节bytes)
        """
        # 处理密钥输入
        if isinstance(key, str):
            key_bytes = bytes.fromhex(key.replace(" ", ""))
        else:
            key_bytes = bytes(key)
            
        # 验证密钥长度
        if len(key_bytes) not in (16, 24, 32):
            raise ValueError("密钥长度必须为16/24/32字节")
            
        # 处理数据输入
        if isinstance(data, str):
            data_bytes = bytes.fromhex(data.replace(" ", ""))
        else:
            data_bytes = bytes(data) if not isinstance(data, bytes) else data
        
        # RFC4493标准填充
        pad_length = 16 - (len(data_bytes) % 16) if len(data_bytes) % 16 != 0 else 0
        if pad_length > 0:
            data_bytes += b'\x80' + b'\x00'*(pad_length-1)
        
        K1, K2 = AESCMAC.generate_subkeys(key_bytes)
        # 初始化状态
        state = bytearray(16)
        
        # 处理数据块
        blocks = [data_bytes[i:i+16] for i in range(0, len(data_bytes), 16)]
        last_block = blocks[-1]
        
        # 前n-1个块处理
        for block in blocks[:-1]:
            state = AESCMAC.xor_bytes(state, block)
            state = AESCMAC.aes_encrypt(state, key_bytes)
        
        # 处理最后一个块（根据RFC4493：未填充时用K1，填充时用K2）
        if pad_length == 0:
            last_block = AESCMAC.xor_bytes(last_block, K1)
        else:
            last_block = AESCMAC.xor_bytes(last_block, K2)
        
        # 处理最后一个块
        state = AESCMAC.xor_bytes(state, last_block)
        cmac = AESCMAC.aes_encrypt(state, key_bytes)
        
        return cmac

    @staticmethod 
    def calculate_hex(data, key = '2b7e1516 28aed2a6 abf71588 09cf4f3c'):
        """
        计算AES-CMAC值(hex字符串形式)
        :param data: 输入数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :return: CMAC结果(32字符hex字符串)
        """
        return AESCMAC.calculate(data, key).hex().upper()

    @staticmethod
    def verify(data, key = '2b7e1516 28aed2a6 abf71588 09cf4f3c', expected_cmac = ''):
        """
        验证AES-CMAC值
        :param data: 输入数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :param expected_cmac: 预期CMAC值(hex字符串或bytes)
        :return: bool 验证结果
        """
        actual_cmac = AESCMAC.calculate(data, key)
        if isinstance(expected_cmac, str):
            expected_bytes = bytes.fromhex(expected_cmac.replace(" ", ""))
        else:
            expected_bytes = bytes(expected_cmac)
        return actual_cmac == expected_bytes

    @staticmethod
    def sha256_hash(data, return_hex=False):
        """
        计算数据的SHA-256哈希值
        :param data: 输入数据(hex字符串或bytes)
        :param return_hex: 是否返回十六进制字符串（默认False返回bytes）
        :return: SHA-256哈希结果(bytes或64字符hex字符串)
        """
        # 处理数据输入
        if isinstance(data, str):
            data_bytes = bytes.fromhex(data.replace(" ", ""))
        else:
            data_bytes = bytes(data) if not isinstance(data, bytes) else data
        
        # 计算SHA-256哈希
        hasher = hashlib.sha256()
        hasher.update(data_bytes)
        hash_result = hasher.digest()
        
        return hash_result.hex().upper() if return_hex else hash_result

    @staticmethod
    def calculate_sigval(data, key = '2b7e1516 28aed2a6 abf71588 09cf4f3c'):
        """
        先计算数据的SHA-256哈希，再用AES-CMAC算法生成签名值
        :param data: 输入原始数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :return: 最终CMAC结果(16字节bytes)
        """
        # 计算SHA-256哈希（返回bytes）
        hash_val = AESCMAC.sha256_hash(data, return_hex=False)
        # 用哈希值计算AES-CMAC
        return AESCMAC.calculate(hash_val, key)

    


# 测试代码
if __name__ == "__main__":
    # 测试用例
    test_vectors = [
        # {
        #     "key": "2b7e1516 28aed2a6 abf71588 09cf4f3c",
        #     "data": "6bc1bee2 2e409f96 e93d7e11 7393172a",
        #     "cmac": "070a16b4 6b4d4144 f79bdd9d d04a287c"
        # },
        # {
        #     "key": "2b7e1516 28aed2a6 abf71588 09cf4f3c",
        #     "data": "6bc1bee2 2e409f96 e93d7e11 7393172a ae2d8a57 1e03ac9c 9eb76fac 45af8e51 30c81c46 a35ce411",
        #     "cmac": "dfa66747 de9ae630 30ca3261 1497c827"
        # },
        # {
        #     "key": "01020304 05060708 09000102 03040506",
        #     "data": "6bc1bee2 2e409f96 e93d7e11 7393172a",
        #     "cmac": "E18AD597 32AE1029 23DF6401 3079C413"
        # },
        # {
        #     "key": "01020304 05060708 09000102 03040506",
        #     "data": "6bc1bee2 2e409f96 e93d7e11 7393172a ae2d8a57 1e03ac9c 9eb76fac 45af8e51 30c81c46 a35ce411",
        #     "cmac": "1861D07B A5E96229 EED61110 B45E5845"
        # },
        {
            "key": "01 02 03 04 05 06 07 08 10 11 12 13 14 15 16 17",
            "data": "e4 4d 02 13 50 49 4e 6f e4 4d 02 13 50 49 4e 6f",
            "cmac": "33669AD014F49CC980260DA442B999B5"      
        },
    ]
    '''
    SHA256 哈希值: D24BC84BE9E99E02A0807AE0FAB9E918F8B2DDF13E9379BBACDE5122A8324B43
    CMAC: 9386DB56257018DD8F72430DE58DB479
    '''
    # cmac = AESCMAC.calculate_sigval("53 30 30 30 30 30 33 35 34 36 30 30 31 00 00 00 53 30 30 30 30 30 33 35 34 33 30 30 31 00 00 00 28 b8 00 10 98 b9 00 10 e8 32 00 20 49 44 01 24 02 20 49 44 03 24 04 21 48 44 05 24 06 21 48 44 07 24 08 22 09 22 47 42 47 43 32 85 00 08 02 09 01 70 71 f1 45 42 02 70 6e f1 01 08 03 09 44 42 20 85 21 85 04 08 06 09 01 70 65 f1 41 44 04 70 62 f8 24 5c ff 03 14 69 5e f1 3f 44 40 70 5b f8 08 08 09 09 01 70 57 f1 3c 44 04 70 54 f0 00 08 02 09 01 70 50 f1 35 42 02 70 4d f1 05 08 07 09 37 42 20 85 21 85 00 08 02 09 01 70 44 f1 2f 42 02 70 41 f1 08 08 09 09 01 70 3d f1 22 5c ff 03 90 d4 39 f1 04 08 06 09 01 70 35 f1 29 44 04 70 32 f8 24 5c ff 03 21 d5 2e f1 27 44 40 70 2b f8 28 42 20 85 32 82 31 70 26 f1 27 42 20 5c fb 07 20 85 32 82 30 70 1f f1 24 42 25 40 20 85 25 44 25 43 13 82 0b 54 04 04 41 30 02 f0 fa 07 40 38 12 f0 04 08 11 10 8a 7a 01 64 04 20 06 20 08 08 01 66 08 20 09 20 c7 f8 00 10 0e 42 0e 43 32 85 00 10 05 04 0c 10 0b 42 0b 43 32 85 00 20 01 20 02 20 03 20 04 20 05 20 06 20 07 20 e8 2a 1d 5c 30 43 65 27 04 03 04 81 de dc 96 78 07 14 15 92 07 00 af 50 00 01 20 40 88 66 5a 5a 28 01 20 40 00 00 08 00 00 02 00 00 00 00 00 00 2c 01 20 40 1c 01 20 40 08 01 20 40 0c 01 20 40 01 00 af 50 88 13 00 00 20 01 20 40 25 5d e0 33 00 20 4f 44 01 24 02 20 4e 44 03 24 04 21 4e 44 05 24 06 21 4d 44 07 24 08 22 09 22 0a 23 0b 23 4b 42 4c 43 32 85 04 08 06 09 01 70 74 f1 4a 44 04 70 71 f8 24 5c 07 00 14 69 6d f1 47 44 40 70 6a f8 08 08 09 09 01 70 66 f1 45 44 04 70 63 f0 00 08 02 09 01 70 5f f1 42 42 02 70 5c f1 01 08 03 09 41 42 20 85 21 85 00 08 02 09 01 70 53 f1 3c 42 02 70 50 f1 05 08 07 09 3c 42 20 85 21 85 00 08 02 09 01 70 47 f1 36 42 02 70 44 f1 04 08 06 09 01 70 40 f1 08 08 09 09 01 70 3c f1 27 10 02 69 39 f1 2b 42 13 4c 32 80 2b 14 32 70 33 f1 2f 42 20 5c f5 07 21 5c f5 07 20 85 32 82 31 70 05 f0 24 42 22 40 20 85 26 04 04 08 06 09 01 70 22 f1 08 09 09 0a 21 70 1e f1 0a 0a 0b 0b 23 70 1a f1 32 72 03 73 32 72 03 73 18 30 21 43 22 44 34 85 22 43 22 45 2a 50 03 04 31 30 07 f0 23 5c 07 00 31 69 08 f1 10 38 ed f1 11 42 11 43 32 85 00 10 05 04 0c 10 0e 42 0e 43 32 85 00 20 01 20 02 20 03 20 04 20 05 20 06 20 07 20 08 20 09 20 0a 20 0b 20 e0 2b 05 5d 1d 5c 00 00 30 43 65 27 04 03 04 81 de dc 96 78 07 14 15 92 07 00 af 50 00 01 20 40 00 00 08 00 00 02 00 00 00 00 00 00 88 66 5a 5a 28 01 20 40 2c 01 20 40 08 01 20 40 0c 01 20 40 01 00 af 50 88 13 00 00 20 01 20 40", key="2b7e1516 28aed2a6 abf71588 09cf4f3c")
    # print(f"计算结果: {cmac.hex().upper()}")
    # print(f"预期结果: 08b07c276fdffd8b5ed6c955894a090b")

    for i, vector in enumerate(test_vectors, 1):
        # cmac = AESCMAC.calculate_sigval(vector['data'],vector['key'])
        # print(f"计算结果: {cmac.hex().upper()}")
        print(f"\n测试用例 {i}:")
        cmac = AESCMAC.calculate_hex(vector["data"], vector["key"])
        print(f"密钥: {vector['key'].replace(' ', '')}")
        print(f"数据: {vector['data'].replace(' ', '')}")
        print(f"计算CMAC: {cmac}")
        print(f"预期CMAC: {vector['cmac'].replace(' ', '').upper()}")
        print(f"验证结果: {'通过' if cmac == vector['cmac'].replace(' ', '').upper() else '失败'}")
        print(f"计算结果: {cmac}")
        print(f"预期结果: {vector['cmac'].replace(' ', '')}")
        print(f"验证结果: {'通过' if AESCMAC.verify(vector['data'], vector['key'], vector['cmac']) else '失败'}")

        
        

